package available

import (
	"context"
	"google.golang.org/grpc"
	"time"
)

type LeakyBucketBuilder struct {
	close chan struct{}
	ch    chan struct{}
}

func (l *LeakyBucketBuilder) Build() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		select {
		case <-l.close:
			return handler(ctx, req)
		case <-l.ch:
			return handler(ctx, req)
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

func (l *LeakyBucketBuilder) Close() {
	close(l.close)
}

func NewLeakyBucketBuilder(duration time.Duration) *LeakyBucketBuilder {
	l := &LeakyBucketBuilder{
		close: make(chan struct{}),
	}
	ticker := time.NewTicker(duration)
	go func() {
		defer ticker.Stop()
		select {
		case <-l.close:
			return
		case <-ticker.C:
			select {
			case l.ch <- struct{}{}:
			default:
			}
		}
	}()

	return l
}
